# SPDX-FileCopyrightText: 2024 German Aerospace Center (DLR.e.V.)
# SPDX-FileContributor: Ferdinand Rewicki <ferdinand.rewicki@dlr.de>
#
# SPDX-License-Identifier: Apache-2.0
"""Methods for finding the ideal number of clusters.

The two types are:
1. scoring function that is maximized to find the ideal number of clusters
2. Methods that determine the value directly
"""
import abc
from abc import abstractmethod
from functools import partial

import aeon
import mlflow
import numpy as np
from sklearn.metrics import adjusted_rand_score, fowlkes_mallows_score
from sklearn.metrics import silhouette_score as sk_silhouette_score
from tslearn.clustering import silhouette_score

from saai.saai_score import saai_score
from saai.tsxmeans import TSXMeans as TSXMeans_


class ScoringFunction(abc.ABC):
    """Abstract scoring function for clustering results."""

    NAME = "NOT_SET"

    @staticmethod
    @abstractmethod
    def scoring_function(X, y_pred, k, **kwargs) -> float:
        """Apply scoring function to clustering result."""
        pass


class KFinder(abc.ABC):
    """Abstract Method for finding K."""

    NAME = "NOT_SET"

    @classmethod
    @abstractmethod
    def find_k(cls, X, clusterer: partial, min_k, max_k, **kwargs) -> float:
        """Determine optimal value for k."""
        pass


class ARI(ScoringFunction):
    """Adjusted Rand Index."""

    NAME = "ari"

    @staticmethod
    def scoring_function(X, y_pred, k, **kwargs) -> float:
        """Calculate Adjusted Rand Index.

        Parameters
        ----------
        X : np-list  The time series data, i.e. anomalous subsequences.
        y_pred : List[int]    The predicted clusters.
        k : The number of clusters K.
        kwargs :
            y_true : List[int]  The ground truth class assignments.

        Returns
        -------
        float
        """
        y_true = kwargs["y_true"]
        return adjusted_rand_score(y_pred, y_true)


class FMI(ScoringFunction):
    """Fowlkes Mallows Index."""

    NAME = "fmi"

    @staticmethod
    def scoring_function(X, y_pred, k, **kwargs) -> float:
        """Calculate Fowlkes-Mallows Index.

        Parameters
        ----------
        X : np-list  The time series data, i.e. anomalous subsequences.
        y_pred : List[int]    The predicted clusters.
        k : The number of clusters K.
        kwargs :
            y_true : List[int]  The ground truth class assignments.

        Returns
        -------
        float
        """
        y_true = kwargs["y_true"]
        return fowlkes_mallows_score(y_pred, y_true)


class SAAI(ScoringFunction):
    """Synchronized Anomaly Agreement Index."""

    NAME = "saai"

    @staticmethod
    def scoring_function(X, y_pred, k, **kwargs) -> float:
        """Calculate SAAI.

        Parameters
        ----------
        X : np-list  The time series data, i.e. anomalous subsequences.
        y_pred : List[int]    The predicted clusters.
        k : int     The number of clusters K.
        kwargs :
            idx_lookup : List[Tuple]    Start and end indices of the anomalous
                subsequences in the original time series

        Returns
        -------
        float
        """
        idx_lookup = kwargs["idx_lookup"]
        score, A_S, A_S_star, n_pseudo_clusters = saai_score(
            idx_lookup, y_pred, k, 0.5, 0.5
        )

        mlflow.log_metric("mt", A_S_star / A_S if A_S > 0 else 0.0, step=k)
        mlflow.log_metric("p1", 1 / k, step=k)
        mlflow.log_metric("p2", n_pseudo_clusters / k, step=k)
        mlflow.log_metric("A_S", A_S, step=k)
        mlflow.log_metric("A_Sc", A_S_star, step=k)
        mlflow.log_metric("n1", n_pseudo_clusters, step=k)

        return score


class TSLearnSSC(ScoringFunction):
    """Silhouette Score (tslearn)."""

    NAME = "ssc"

    @staticmethod
    def scoring_function(X, y_pred, k, **kwargs) -> float:
        """Calculate Silhouette Score.

        Calculate Silhouette Score using the tslearn implementation

        Parameters
        ----------
        X : np-list  The time series data, i.e. anomalous subsequences.
        y_pred : List[int]    The predicted clusters.
        k : The number of clusters K.
        kwargs : -

        Returns
        -------
        float
        """
        return silhouette_score(X.squeeze(), y_pred, metric="dtw")


class SKLearnSSC(ScoringFunction):
    """Silhouette Score (scikit-learn)."""

    NAME = "skssc"

    @staticmethod
    def scoring_function(X, y_pred, k, **kwargs) -> float:
        """Calculate Silhouette Score.

        Calculate Silhouette Score using the scikit-learn implementation

        Parameters
        ----------
        X : np-list  The time series data, i.e. anomalous subsequences.
        y_pred : List[int]    The predicted clusters.
        k : The number of clusters K.
        kwargs : -

        Returns
        -------
        float
        """
        return sk_silhouette_score(
            [np.nan_to_num(x.squeeze()) for x in X],
            y_pred,
            metric=aeon.distances.dtw_distance,
        )


class TSXMeans(KFinder):
    """Time Series XMeans."""

    NAME = "tsxm"

    @classmethod
    def find_k(cls, X, max_k, **kwargs) -> float:
        """Find optimal k using TSXMeans.

        Parameters
        ----------
        X : np-list The time series data, i.e. anomalous subsequences.
        max_k : integer Maximal value for K
        kwargs : -

        Returns
        -------
        integer Optimal K
        """
        xmeans = TSXMeans_(kmax=max_k, max_iter=50, distance="msm")
        xmeans.fit(np.nan_to_num(X))

        mlflow.log_metric("tsxm_k", xmeans.n_clusters)

        return xmeans.n_clusters
